#include "StdAfx.h"
#include "WinsockHooks.h"
#include <detours.h>
#include <tchar.h>

#include "NetworkEvent.h"
#include "PostOffice.h"
#include "DetouredFunctions.hpp"

/*static*/ WinsockHooks * WinsockHooks::_instance = NULL;
/*static*/ boost::once_flag WinsockHooks::_instance_flag;


/**
 * Constructor. Initializes list of regex'es to identify
 * "interesting" DLLs, as described in the comment in H file.
 * TODO: read this list from configuration file/registry
 */
WinsockHooks::WinsockHooks(void)
	: _is_activated(false), _hm_service_provider(NULL), original_WSPStartup(NULL),
	   original_gethostbyname(NULL)
{
	_service_providers_dlls.push_back(tregex(_T(".*MSWSOCK.*"), boost::regex::perl | boost::regex::icase));
}

void WinsockHooks::doProcessDLL(const platform::XString &name, HMODULE hm) {
	// check whether is one of known DLLs
	bool known = false;
	for (ListRegex::iterator it = _service_providers_dlls.begin(); it != _service_providers_dlls.end(); ++it)
		if (boost::regex_match(name.getTStr(), *it)) {
			known = true;
			break;
		}
	if (!known)
		// do nothing for this DLL
		return;

	// need to hook the WSPStartup function
	// but let's first make sure that it exists
	FARPROC fStartup = ::GetProcAddress(hm, "WSPStartup");
	if (fStartup == NULL)
		// nothing to do with this function
		return;
	_hm_service_provider = hm;
	original_WSPStartup = (WSPAPI_WSPStartup) fStartup;

	// find "gethostbyname" function
	HMODULE hmWinsock = ::LoadLibrary(_T("Ws2_32.dll"));
	if (hmWinsock != NULL)
		original_gethostbyname = (WINSOCK_gethostbyname) ::GetProcAddress(hmWinsock, "gethostbyname");

	// now let's install the hooks
	DetourTransactionBegin();
	DetourUpdateThread(::GetCurrentThread());
	DetourAttach(&(PVOID&)original_WSPStartup, HookWSPStartup);
	if (original_gethostbyname != NULL)
		DetourAttach(&(PVOID&)original_gethostbyname, TrapGethostbyname);
	DetourTransactionCommit();
}


/**
 * Replaces original function pointers with our versions
 */
/*static*/ int WSPAPI WinsockHooks::HookWSPStartup(
	WORD wVersionRequested,
	LPWSPDATA lpWSPData,
	LPWSAPROTOCOL_INFOW lpProtocolInfo,
	WSPUPCALLTABLE UpcallTable,
	LPWSPPROC_TABLE lpProcTable
)
{
	WinsockHooks &inst = GetInstance();

	int res = inst.original_WSPStartup(wVersionRequested, lpWSPData, lpProtocolInfo, UpcallTable, lpProcTable);
	if (res == 0) {
		// store original functions
		memcpy(&inst.original_WSPFunctions, lpProcTable, sizeof(inst.original_WSPFunctions));

		// replace some of the functions with hooked version
		lpProcTable->lpWSPConnect	= TrapWSPConnect;
		lpProcTable->lpWSPSend		= TrapWSPSend;
		lpProcTable->lpWSPRecv		= TrapWSPRecv;
		lpProcTable->lpWSPSocket	= TrapWSPSocket;
		lpProcTable->lpWSPListen	= TrapWSPListen;
		lpProcTable->lpWSPSendTo	= TrapWSPSendTo;
		lpProcTable->lpWSPRecvFrom	= TrapWSPResvFrom;
	}
	return res;
}


/** LPWSPCONNECT */
/*static*/ int
WSPAPI
WinsockHooks::TrapWSPConnect(
	SOCKET s,
	const struct sockaddr FAR * name,
	int namelen,
	LPWSABUF lpCallerData,
	LPWSABUF lpCalleeData,
	LPQOS lpSQOS,
	LPQOS lpGQOS,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_CONNECT, s);
		if (name != NULL && name->sa_family == AF_INET) {
			struct sockaddr_in * sin = (struct sockaddr_in *)name;
			struct in_addr &addr = sin->sin_addr;
			ne->setIp(addr.S_un.S_addr);
		}
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPConnect(s, name, namelen, lpCallerData, lpCalleeData, lpSQOS, lpGQOS, lpErrno);
}

/** LPWSPSEND */
/*static*/ int
WSPAPI
WinsockHooks::TrapWSPSend(
	SOCKET s,
	LPWSABUF lpBuffers,
	DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesSent,
	DWORD dwFlags,
	LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
	LPWSATHREADID lpThreadId,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_SEND, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPSend(s, lpBuffers, dwBufferCount, lpNumberOfBytesSent, dwFlags, lpOverlapped, lpCompletionRoutine, lpThreadId, lpErrno);
}

/** LPWSPRECV */
/*static*/ int
WSPAPI
WinsockHooks::TrapWSPRecv(
	SOCKET s,
	LPWSABUF lpBuffers,
	DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesRecvd,
	LPDWORD lpFlags,
	LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
	LPWSATHREADID lpThreadId,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_RECEIVE, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPRecv(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpOverlapped, lpCompletionRoutine, lpThreadId, lpErrno);
}

/** LPWSPSOCKET */
/*static*/ SOCKET
WSPAPI
WinsockHooks::TrapWSPSocket(
	int af,
	int type,
	int protocol,
	LPWSAPROTOCOL_INFOW lpProtocolInfo,
	GROUP g,
	DWORD dwFlags,
	LPINT lpErrno
	)
{
	SOCKET s = WinsockHooks::GetInstance().original_WSPFunctions.lpWSPSocket(af, type, protocol, lpProtocolInfo, g, dwFlags, lpErrno);
	ErrorCodeRestore erc;
	{
		StateIndicator::Setter isInternal;
		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_CREATE_SOCKET, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return s;
}

/** LPWSPLISTEN */
/*static*/ int
WSPAPI
WinsockHooks::TrapWSPListen(
	SOCKET s,
	int backlog,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_LISTEN, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPListen(s, backlog, lpErrno);
}

/** LPWSPSENDTO */
/*static*/ int
WSPAPI
WinsockHooks::TrapWSPSendTo(
	SOCKET s,
	LPWSABUF lpBuffers,
	DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesSent,
	DWORD dwFlags,
	const struct sockaddr FAR * lpTo,
	int iTolen,
	LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
	LPWSATHREADID lpThreadId,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_SEND_TO, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPSendTo(s, lpBuffers, dwBufferCount, lpNumberOfBytesSent, dwFlags, lpTo, iTolen, lpOverlapped, lpCompletionRoutine, lpThreadId, lpErrno);
}

/** LPWSPRECVFROM */
/*static*/ int WSPAPI
WinsockHooks::TrapWSPResvFrom(
	SOCKET s,
	LPWSABUF lpBuffers,
	DWORD dwBufferCount,
	LPDWORD lpNumberOfBytesRecvd,
	LPDWORD lpFlags,
	struct sockaddr FAR * lpFrom,
	LPINT lpFromlen,
	LPWSAOVERLAPPED lpOverlapped,
	LPWSAOVERLAPPED_COMPLETION_ROUTINE lpCompletionRoutine,
	LPWSATHREADID lpThreadId,
	LPINT lpErrno
	)
{
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_RECEIVE_FROM, s);
		events::PostOffice::GetInstance().send(ne);
	}
	return WinsockHooks::GetInstance().original_WSPFunctions.lpWSPRecvFrom(s, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, lpFlags, lpFrom, lpFromlen, lpOverlapped, lpCompletionRoutine, lpThreadId, lpErrno);
}

/*static*/ struct hostent * PASCAL
WinsockHooks::TrapGethostbyname(const char * name)
{
	/*
	typedef struct hostent {
		char FAR *h_name;
		char FAR  FAR **h_aliases;
		short h_addrtype;
		short h_length;
		char FAR  FAR **h_addr_list;
	} HOSTENT;
	*/
	struct hostent * pret = WinsockHooks::GetInstance().original_gethostbyname(name);
	ErrorCodeRestore erc;
	{
		StateIndicator::Setter isInternal;

		events::NetworkEvent *ne = new events::NetworkEvent(events::NetworkEvent::NE_GET_HOST_BY_NAME, 0);
		ne->setHost(std::string(name == NULL ? "" : name));
		if (pret != NULL)
			if (pret->h_addrtype == AF_INET && pret->h_length == sizeof(unsigned long))
				ne->setIp(*((unsigned long*)pret->h_addr_list[0]));
		events::PostOffice::GetInstance().send(ne);
	}	
	return pret;
}

